package com.lhkk.inventoryvehicle.ai


import android.content.Context
import android.graphics.*
import android.util.AttributeSet
import android.view.View
import org.tensorflow.lite.task.vision.detector.Detection

class OverlayView(context: Context?, attrs: AttributeSet?) : View(context, attrs) {

    private var results: List<Detection> = listOf()
    private var imageWidth: Int = 1
    private var imageHeight: Int = 1

    private val boxPaint = Paint().apply {
        color = Color.GREEN
        style = Paint.Style.STROKE
        strokeWidth = 8f
    }

    private val textPaint = Paint().apply {
        color = Color.WHITE
        style = Paint.Style.FILL
        textSize = 50f
    }

    // 虚拟触发线
    private val linePaint = Paint().apply {
        color = Color.YELLOW
        style = Paint.Style.STROKE
        strokeWidth = 6f
        pathEffect = DashPathEffect(floatArrayOf(20f, 20f), 0f)
    }

    override fun onDraw(canvas: Canvas) {
        super.onDraw(canvas)

        // 绘制虚拟线
        // 线A: 位于屏幕 1/3 处
        val lineA_X = width * 0.33f
        canvas.drawLine(lineA_X, 0f, lineA_X, height.toFloat(), linePaint)
        // 线B: 位于屏幕 2/3 处
        val lineB_X = width * 0.66f
        canvas.drawLine(lineB_X, 0f, lineB_X, height.toFloat(), linePaint)

        // 绘制检测到的车辆边界框
        for (result in results) {
            val boundingBox = result.boundingBox
            val scaleX = width.toFloat() / imageWidth
            val scaleY = height.toFloat() / imageHeight

            val scaledBox = RectF(
                boundingBox.left * scaleX,
                boundingBox.top * scaleY,
                boundingBox.right * scaleX,
                boundingBox.bottom * scaleY
            )
            canvas.drawRect(scaledBox, boxPaint)

            val label = "${result.categories.first().label} %.2f".format(result.categories.first().score)
            canvas.drawText(label, scaledBox.left, scaledBox.top - 10, textPaint)
        }
    }

    fun setResults(detections: List<Detection>, imageWidth: Int, imageHeight: Int) {
        this.results = detections.filter {
            // 只显示车辆相关的检测结果
            it.categories.first().label in listOf("car", "truck")
        }
        this.imageWidth = imageWidth
        this.imageHeight = imageHeight
        invalidate()
    }
}